import torch
import numpy as np
import torch.nn as nn
import torch.nn.functional as F


class TestNet(nn.Module):
    def __init__(self):
        super(TestNet, self).__init__()

    def forward(self,input_x, scale, bias, mean, variance):
        out = F.batch_norm(input_x, mean, variance, scale, bias)
        return [out]


if __name__ == "__main__":
    net = TestNet()

    # x = torch.FloatTensor(np.ones([2,2]).astype(np.float32))
    # mean = torch.FloatTensor(np.ones([2]).astype(np.float32))
    # variance = torch.FloatTensor(np.ones([2]).astype(np.float32))
    # scale = torch.FloatTensor(np.ones([2]).astype(np.float32))
    # bias = torch.FloatTensor(np.ones([2]).astype(np.float32))

    x = torch.FloatTensor([[0.9253564,1.2399976], [-1.444724,-0.3179193]])
    scale = torch.FloatTensor([-2.3426611,  0.9450944])
    bias = torch.FloatTensor([0.7408705, 2.0633006])
    mean = torch.FloatTensor([0.5742812, 1.0141671])
    variance = torch.FloatTensor([ 0.34994918, -1.1426003 ])
    result = net(x, scale, bias, mean, variance)
    print(result[0])
    print(result[0].mean())
    print(result[0].shape)


